#!/bin/bash
set -eo pipefail
set -E

# 作者：无涯有涯
# 功能：恢复CloudNativePG管理的PostgreSQL集群数据库
# 用法：curl -sSL https://gitee.com/veryx/galaxy-tools/raw/master/5-cnpg/pg_restore.sh | bash -s -- -d <db> -f <file> -u <user>

# 参数
DB_NAME=""
DB_USER=""
BACKUP_FILE=""

# 日志
LOG_FILE="${LOG_FILE:-$HOME/.kapp/log/pg_restore_$(date +%Y%m%d-%H%M%S).log}"

begin=$(date +%s)

# 显示用法
usage() {
    echo "用法: $(basename "$0") -u <用户名> -d <数据库名> -f <备份文件>"
    echo "示例: $(basename "$0") -u cosmic -d demo_full -f /tmp/baseline_show.backup"
}

# 参数检查
if [[ $# -eq 0 ]]; then
    usage
    exit 0
fi

while [[ $# -gt 0 ]]; do
    case "$1" in
        -h|--help) usage; exit 0 ;;
        -u) DB_USER="$2"; shift 2 ;;
        -d) DB_NAME="$2"; shift 2 ;;
        -f) BACKUP_FILE="$2"; shift 2 ;;
        *) echo "错误: 未知参数 $1"; exit 1 ;;
    esac
done

# 验证必需参数
if [[ -z "$DB_USER" || -z "$DB_NAME" || -z "$BACKUP_FILE" ]]; then
    echo "错误: 缺少必填参数！"
    usage
    exit 1
fi

# 验证数据库名格式
if [[ ! "$DB_NAME" =~ ^[a-z][a-z0-9_]{0,62}$ ]]; then
    echo "错误: 数据库名需为 1-63 个字符，只能包含小写字母、数字和下划线，且必须以字母开头。"
    exit 1
fi

# 验证备份文件
if [ ! -f "$BACKUP_FILE" ]; then
    echo "错误: 文件不存在: $BACKUP_FILE"
    exit 1
fi

if [ ! -r "$BACKUP_FILE" ]; then
    echo "错误: 文件不可读: $BACKUP_FILE"
    echo "建议: sudo chmod +r \"$BACKUP_FILE\" 为文件添加读取权限。"
    exit 1
fi

# 获取绝对路径
BACKUP_FILE=$(readlink -f "$BACKUP_FILE")

# 获取PostgreSQL主节点
POD=$(kubectl get pod -n pg-system -l 'cnpg.io/instanceRole=primary' -o jsonpath='{.items[0].metadata.name}' 2>/dev/null)
if [ -z "$POD" ]; then
    echo "错误: 无法获取PostgreSQL集群的主节点信息"
    exit 1
fi

# 创建日志目录
mkdir -p "$(dirname "$LOG_FILE")"

# 检查数据库是否存在
DB_EXISTS=$(kubectl exec -n pg-system "$POD" -c postgres -- psql -d postgres -tc "SELECT 1 FROM pg_database WHERE datname='$DB_NAME'" | tr -d '[:space:]')

if [ "$DB_EXISTS" = "1" ]; then
    echo -ne "\e[33m警告：\e[0m\n数据库 '$DB_NAME' 已存在！继续操作将导致现有数据丢失。是否继续？(y/N): "
    read confirm
    if [[ ! "$confirm" =~ ^[yY]$ ]]; then
        echo "操作已取消"
        exit 0
    fi
    echo "目标数据库：$DB_NAME，源文件：$BACKUP_FILE，已存在同名数据库，用户选择\"$confirm\"覆盖现有数据库。" | tee -a "$LOG_FILE"
else
    echo "目标数据库：$DB_NAME，源文件：$BACKUP_FILE，不存在同名数据库，将创建新数据库。" | tee -a "$LOG_FILE"
fi

# 容器内文件路径
TMP_BAK_FILE="/var/lib/postgresql/data/$(basename "$BACKUP_FILE")"
TMP_TOC_FILE="/var/lib/postgresql/data/$(basename "$DB_NAME").toc"

# 传输备份文件
echo "传输备份文件中..."
kubectl cp "$BACKUP_FILE" "pg-system/$POD:$TMP_BAK_FILE" --retries=-1 -c postgres || {
    echo "错误: 备份文件复制到容器失败！"
    exit 1
}

# 执行数据库恢复
echo "开始恢复数据库... 执行日志: $LOG_FILE"
echo "数据库恢复正在进行中，请勿中断进程..."

(kubectl exec -n pg-system -i "$POD" -c postgres -- bash -e <<EOF
set -eo pipefail
set -E

cd /var/lib/postgresql/data || exit 1

# 创建数据库
psql -d postgres -tc "SELECT 1 FROM pg_database WHERE datname='$DB_NAME'" | grep -q 1 || createdb -E UTF8 -O "$DB_USER" "$DB_NAME"

# 生成恢复清单
pg_restore -l "$TMP_BAK_FILE" | grep -Ev 'EXTENSION|ACL|rdsAdmin|spatial_ref_sys|pg_stat_statements|geometry_columns|geography' > "$TMP_TOC_FILE"

# 执行备份恢复
pg_restore -j 8 -v -d "$DB_NAME" -c --if-exists -L "$TMP_TOC_FILE" --no-owner --role="$DB_USER"  "$TMP_BAK_FILE"

# 清理临时文件
trap 'rm -f "$TMP_BAK_FILE" "$TMP_TOC_FILE"' INT HUP EXIT TERM QUIT ABRT
EOF
) >> "$LOG_FILE" 2>&1 || {
    echo "错误: 数据库恢复失败！"
    exit 1
}

duration=$(( $(date +%s) - begin ))
echo "数据库恢复完成。耗时：$((duration / 60)) 分 $((duration % 60)) 秒" | tee -a "$LOG_FILE"